{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyNw3Y7vOkMMxSnIQMGGEMMc",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "567751bf6abc4f7b97755ff328543f45": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_ec058cd9ef4444b1ad303496cab7db2f",
              "IPY_MODEL_7184df2298a840dab7c9c137b3a9f245",
              "IPY_MODEL_9735baec166944409934152b369f4387"
            ],
            "layout": "IPY_MODEL_87a2593743594c8d8ae32a0768e7e1c1"
          }
        },
        "ec058cd9ef4444b1ad303496cab7db2f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_aee96ba3886f4d208a82024e56fa75e3",
            "placeholder": "​",
            "style": "IPY_MODEL_37ebb56bf4b24fa5a2ec6ce66c1b8d08",
            "value": "Map: 100%"
          }
        },
        "7184df2298a840dab7c9c137b3a9f245": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_3859a1caa9c244d4957f40704423e2f5",
            "max": 366,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_29887aaeb71240559ceca8babd13c0a3",
            "value": 366
          }
        },
        "9735baec166944409934152b369f4387": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_761da1e6531e48f59d911320075410b0",
            "placeholder": "​",
            "style": "IPY_MODEL_93df3ac3106d43a6b71e42cb8728b841",
            "value": " 366/366 [00:00&lt;00:00, 2910.54 examples/s]"
          }
        },
        "87a2593743594c8d8ae32a0768e7e1c1": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "aee96ba3886f4d208a82024e56fa75e3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "37ebb56bf4b24fa5a2ec6ce66c1b8d08": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "3859a1caa9c244d4957f40704423e2f5": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "29887aaeb71240559ceca8babd13c0a3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "761da1e6531e48f59d911320075410b0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "93df3ac3106d43a6b71e42cb8728b841": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "10af351daecf47c7bc43273881506790": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_b27846081b5a4ebe8dc4451e89f3f887",
              "IPY_MODEL_11ebb8f4fd0041468c4a9597a1c7622a",
              "IPY_MODEL_3fa5b79dcba0443986f2093a60503675"
            ],
            "layout": "IPY_MODEL_2bc9f8fb52d94817a789c213ec73bc01"
          }
        },
        "b27846081b5a4ebe8dc4451e89f3f887": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_019813bb90034402a6b34ac776777a5d",
            "placeholder": "​",
            "style": "IPY_MODEL_7e86fdb7f8f748699af6059535a18f7c",
            "value": "Loading checkpoint shards: 100%"
          }
        },
        "11ebb8f4fd0041468c4a9597a1c7622a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_17013bed677942cdb0e28861e0a9be4b",
            "max": 2,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_7b49e9a7b12041a38564eaf51c054e85",
            "value": 2
          }
        },
        "3fa5b79dcba0443986f2093a60503675": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7f914a658d58492d97f6c9fd0adddadb",
            "placeholder": "​",
            "style": "IPY_MODEL_361ec8f8ae4b4d48bab82130e2064da8",
            "value": " 2/2 [00:59&lt;00:00, 27.02s/it]"
          }
        },
        "2bc9f8fb52d94817a789c213ec73bc01": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "019813bb90034402a6b34ac776777a5d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7e86fdb7f8f748699af6059535a18f7c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "17013bed677942cdb0e28861e0a9be4b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7b49e9a7b12041a38564eaf51c054e85": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "7f914a658d58492d97f6c9fd0adddadb": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "361ec8f8ae4b4d48bab82130e2064da8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "a593780524b4420da3cc4e907bcc8b07": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_ce9323cc706e4752bc9474c23daba43a",
              "IPY_MODEL_17a2a78957654181a21d7bf86d6caaaf",
              "IPY_MODEL_09299c4279be41ea932756b77380ca58"
            ],
            "layout": "IPY_MODEL_c715443800cd4a129ad4741116363a2f"
          }
        },
        "ce9323cc706e4752bc9474c23daba43a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_af0147dc68a245d7ad01c088dc21ab7c",
            "placeholder": "​",
            "style": "IPY_MODEL_265014e8078b4484ac590d5ea10c0804",
            "value": "Map: 100%"
          }
        },
        "17a2a78957654181a21d7bf86d6caaaf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b97668c0895346f29d6f75a0e77175e9",
            "max": 366,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_cbe5954851214e2a91ce943adf0c64d5",
            "value": 366
          }
        },
        "09299c4279be41ea932756b77380ca58": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ce3dfb1882e646d0ac247eda73cd2565",
            "placeholder": "​",
            "style": "IPY_MODEL_2ec5b9ac76a54e83b2fb9a4ec319acd1",
            "value": " 366/366 [00:00&lt;00:00, 979.94 examples/s]"
          }
        },
        "c715443800cd4a129ad4741116363a2f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "af0147dc68a245d7ad01c088dc21ab7c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "265014e8078b4484ac590d5ea10c0804": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b97668c0895346f29d6f75a0e77175e9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "cbe5954851214e2a91ce943adf0c64d5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "ce3dfb1882e646d0ac247eda73cd2565": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "2ec5b9ac76a54e83b2fb9a4ec319acd1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d62f31bc09054d72846e601dcdb167f2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_78a776d88c944c558aa8654e08ccae39",
              "IPY_MODEL_eafe0b3a7b324bdbafc6d28da64885dc",
              "IPY_MODEL_7b6253d809ea4b3588800f87b2ddd641"
            ],
            "layout": "IPY_MODEL_9998068daf81422aae5b0e2e3223b11a"
          }
        },
        "78a776d88c944c558aa8654e08ccae39": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4551d7c16dd540eca3d8fe819e14b221",
            "placeholder": "​",
            "style": "IPY_MODEL_62da5e0db3dd47c786d56c7d7828f38b",
            "value": "Loading checkpoint shards: 100%"
          }
        },
        "eafe0b3a7b324bdbafc6d28da64885dc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ec8913ad1253489bafcfdba52fbcb528",
            "max": 2,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a2b8a05d722c40058cc7f6868d9f086b",
            "value": 2
          }
        },
        "7b6253d809ea4b3588800f87b2ddd641": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7f14dd60d1aa45fabecf5d98168f7be4",
            "placeholder": "​",
            "style": "IPY_MODEL_4fc60b58df5648cba3f9a8fc865bfeb8",
            "value": " 2/2 [00:54&lt;00:00, 25.13s/it]"
          }
        },
        "9998068daf81422aae5b0e2e3223b11a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4551d7c16dd540eca3d8fe819e14b221": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "62da5e0db3dd47c786d56c7d7828f38b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "ec8913ad1253489bafcfdba52fbcb528": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a2b8a05d722c40058cc7f6868d9f086b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "7f14dd60d1aa45fabecf5d98168f7be4": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4fc60b58df5648cba3f9a8fc865bfeb8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/hengjiUSTC/learn-llm/blob/main/learn_qlora_finetune.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q -U bitsandbytes\n",
        "!pip install -q -U git+https://github.com/huggingface/transformers.git\n",
        "!pip install -q -U git+https://github.com/huggingface/peft.git\n",
        "!pip install -q -U git+https://github.com/huggingface/accelerate.git\n",
        "!pip install -q datasets trl sentencepiece protobuf"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GfMoF97yM78i",
        "outputId": "e70384a1-f592-4dfb-b2c3-31ed110ced90"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prepare data\n",
        "\n",
        "- load 366 rows of Chinese poetry\n",
        "- format data to conform llama2 chat."
      ],
      "metadata": {
        "id": "Em5y8iPOkmYz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "import requests\n",
        "from datasets import Dataset\n",
        "\n",
        "def load_json_from_url(url):\n",
        "    response = requests.get(url)\n",
        "    if response.status_code == 200:\n",
        "        return response.json()\n",
        "    else:\n",
        "        raise Exception(\"Failed to load data from URL\")\n",
        "\n",
        "def generate_prompt(data_point):\n",
        "    tags = ';'.join(data_point['tags'])\n",
        "    paragraph = '\\n'.join(data_point['paragraphs'])\n",
        "    return f\"\"\"[INST] <<SYS>>你是一个唐诗助手,帮助用户写一首对应要求的唐诗<</SYS>>\n",
        "作者:{data_point[\"author\"]}\n",
        "标签:{tags}\n",
        "[/INST]{data_point['title']}\n",
        "{paragraph}\n",
        "\"\"\".strip()\n",
        "\n",
        "def generate_text(data_point):\n",
        "    full_prompt = generate_prompt(data_point)\n",
        "    return {\"text\": full_prompt}\n",
        "\n",
        "# URL of the JSON file\n",
        "url = \"https://raw.githubusercontent.com/chinese-poetry/chinese-poetry/master/%E5%85%A8%E5%94%90%E8%AF%97/%E5%94%90%E8%AF%97%E4%B8%89%E7%99%BE%E9%A6%96.json\"\n",
        "data = load_json_from_url(url)\n",
        "\n",
        "# Convert data to Pandas DataFrame\n",
        "df = pd.DataFrame(data=data)\n",
        "\n",
        "# Create Dataset and apply transformations\n",
        "dataset = Dataset.from_pandas(df)\n",
        "dataset = dataset.shuffle().map(generate_text)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "567751bf6abc4f7b97755ff328543f45",
            "ec058cd9ef4444b1ad303496cab7db2f",
            "7184df2298a840dab7c9c137b3a9f245",
            "9735baec166944409934152b369f4387",
            "87a2593743594c8d8ae32a0768e7e1c1",
            "aee96ba3886f4d208a82024e56fa75e3",
            "37ebb56bf4b24fa5a2ec6ce66c1b8d08",
            "3859a1caa9c244d4957f40704423e2f5",
            "29887aaeb71240559ceca8babd13c0a3",
            "761da1e6531e48f59d911320075410b0",
            "93df3ac3106d43a6b71e42cb8728b841"
          ]
        },
        "id": "NVrvVpB_M1Gu",
        "outputId": "4969781b-b341-4764-b298-a72630c4e928"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Map:   0%|          | 0/366 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "567751bf6abc4f7b97755ff328543f45"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(dataset[0]['text'])\n",
        "dataset"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Wh25Hy6KNhTq",
        "outputId": "7d8e871b-c1e8-4695-dc13-145964cca265"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[INST] <<SYS>>你是一个唐诗助手,帮助用户写一首对应要求的唐诗<</SYS>>\n",
            "作者:劉禹錫\n",
            "标签:唐诗三百首;怀古;七言绝句;带有地名;地名\n",
            "[/INST]金陵五題 烏衣巷 \n",
            "朱雀橋邊野草花，烏衣巷口夕陽斜。\n",
            "舊時王謝堂前燕，飛入尋常百姓家。\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Dataset({\n",
              "    features: ['author', 'paragraphs', 'tags', 'title', 'id', 'text'],\n",
              "    num_rows: 366\n",
              "})"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "ArYMH1eKL5qy"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer\n",
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Parameters for finetune process"
      ],
      "metadata": {
        "id": "j0dzmaE0lHee"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model_id = \"NousResearch/Llama-2-7b-chat-hf\"\n",
        "\n",
        "################################################################################\n",
        "# QLoRA parameters\n",
        "################################################################################\n",
        "\n",
        "# LoRA attention dimension\n",
        "lora_r = 64\n",
        "\n",
        "# Alpha parameter for LoRA scaling\n",
        "lora_alpha = 16\n",
        "\n",
        "# Dropout probability for LoRA layers\n",
        "lora_dropout = 0.1\n",
        "\n",
        "lora_target_modules = [\n",
        "    \"q_proj\",\n",
        "    \"up_proj\",\n",
        "    \"o_proj\",\n",
        "    \"k_proj\",\n",
        "    \"down_proj\",\n",
        "    \"gate_proj\",\n",
        "    \"v_proj\",\n",
        "]\n",
        "\n",
        "################################################################################\n",
        "# bitsandbytes parameters\n",
        "################################################################################\n",
        "\n",
        "# Activate 4-bit precision base model loading\n",
        "use_4bit = True\n",
        "\n",
        "# Compute dtype for 4-bit base models\n",
        "bnb_4bit_compute_dtype = torch.float16\n",
        "\n",
        "# Quantization type (fp4 or nf4)\n",
        "bnb_4bit_quant_type = \"nf4\"\n",
        "\n",
        "# Activate nested quantization for 4-bit base models (double quantization)\n",
        "use_nested_quant = True\n",
        "\n",
        "################################################################################\n",
        "# TrainingArguments parameters\n",
        "################################################################################\n",
        "\n",
        "# Output directory where the model predictions and checkpoints will be stored\n",
        "output_dir = \"results\"\n",
        "save_dir = \"qlora-result\"\n",
        "\n",
        "# Number of training epochs\n",
        "num_train_epochs = 3\n",
        "\n",
        "# Enable fp16/bf16 training (set bf16 to True with an A100)\n",
        "fp16 = True\n",
        "bf16 = False\n",
        "\n",
        "# Batch size per GPU for training\n",
        "per_device_train_batch_size = 1\n",
        "\n",
        "# Number of update steps to accumulate the gradients for\n",
        "gradient_accumulation_steps = 3\n",
        "\n",
        "# Maximum gradient normal (gradient clipping)\n",
        "max_grad_norm = 0.3\n",
        "\n",
        "# Initial learning rate (AdamW optimizer)\n",
        "learning_rate = 2e-4\n",
        "\n",
        "# Weight decay to apply to all layers except bias/LayerNorm weights\n",
        "weight_decay = 0.001\n",
        "\n",
        "# Optimizer to use\n",
        "optim = \"paged_adamw_32bit\"\n",
        "\n",
        "# Learning rate schedule\n",
        "lr_scheduler_type = \"cosine\"\n",
        "\n",
        "# Ratio of steps for a linear warmup (from 0 to learning rate)\n",
        "warmup_ratio = 0.03\n",
        "\n",
        "# Log every X updates steps\n",
        "logging_steps = 10\n",
        "\n",
        "################################################################################\n",
        "# SFT parameters\n",
        "################################################################################\n",
        "\n",
        "# Maximum sequence length to use\n",
        "max_seq_length = 2048\n",
        "\n",
        "# Pack multiple short examples in the same input sequence to increase efficiency\n",
        "packing = False\n",
        "\n",
        "# Load the entire model on the GPU 0\n",
        "device_map = {\"\": 0}"
      ],
      "metadata": {
        "id": "VP42CtF7MBx_"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Load model\n",
        "\n",
        "- load llama2 chat model in int4 format (quantization)"
      ],
      "metadata": {
        "id": "CfnDbftRlNVU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=use_4bit,\n",
        "    bnb_4bit_use_double_quant=use_nested_quant,\n",
        "    bnb_4bit_quant_type=bnb_4bit_quant_type,\n",
        "    bnb_4bit_compute_dtype=bnb_4bit_compute_dtype\n",
        ")\n",
        "\n",
        "\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_id,\n",
        "    quantization_config=bnb_config,\n",
        "    device_map=device_map,\n",
        "    torch_dtype=torch.float16,\n",
        "    use_cache=False\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 138,
          "referenced_widgets": [
            "10af351daecf47c7bc43273881506790",
            "b27846081b5a4ebe8dc4451e89f3f887",
            "11ebb8f4fd0041468c4a9597a1c7622a",
            "3fa5b79dcba0443986f2093a60503675",
            "2bc9f8fb52d94817a789c213ec73bc01",
            "019813bb90034402a6b34ac776777a5d",
            "7e86fdb7f8f748699af6059535a18f7c",
            "17013bed677942cdb0e28861e0a9be4b",
            "7b49e9a7b12041a38564eaf51c054e85",
            "7f914a658d58492d97f6c9fd0adddadb",
            "361ec8f8ae4b4d48bab82130e2064da8"
          ]
        },
        "id": "KUVd4UHlOcbK",
        "outputId": "6ffede9f-1908-460f-ff00-bcedba353d82"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "10af351daecf47c7bc43273881506790"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/transformers/generation/configuration_utils.py:389: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.9` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`. This was detected when initializing the generation config instance, which means the corresponding file may hold incorrect parameterization and should be fixed.\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.10/dist-packages/transformers/generation/configuration_utils.py:394: UserWarning: `do_sample` is set to `False`. However, `top_p` is set to `0.6` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_p`. This was detected when initializing the generation config instance, which means the corresponding file may hold incorrect parameterization and should be fixed.\n",
            "  warnings.warn(\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Build tokenizer\n",
        "\n",
        "- cautious I added special token handling `tokenizer.pad_token_id = 18610`. Without this change finetuning process will cause none stopping issue for generated model."
      ],
      "metadata": {
        "id": "1Nie2R4dlaz9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(model_id,\n",
        "                                          trust_remote_code=True,\n",
        "                                          add_eos_token=True,\n",
        "                                          use_fast=False)\n",
        "tokenizer.add_special_tokens({\n",
        "    \"eos_token\": tokenizer.convert_ids_to_tokens(model.config.eos_token_id),\n",
        "    \"bos_token\": tokenizer.convert_ids_to_tokens(model.config.bos_token_id),\n",
        "    \"unk_token\": tokenizer.convert_ids_to_tokens(\n",
        "        model.config.pad_token_id if model.config.pad_token_id != -1 else tokenizer.pad_token_id\n",
        "    ),\n",
        "})\n",
        "#ref:\n",
        "# https://github.com/huggingface/transformers/issues/22794#issuecomment-1616258519\n",
        "# https://www.reddit.com/r/LocalLLaMA/comments/15hz7gl/my_finetuning_based_on_llama27bchathf_model/\n",
        "tokenizer.pad_token_id = 18610 #_***\n",
        "tokenizer.padding_side = \"right\"\n",
        "tokenizer"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "up21cFwqOe3h",
        "outputId": "17de0767-6ad0-4e99-ba83-a0b0f2feecea"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "LlamaTokenizer(name_or_path='NousResearch/Llama-2-7b-chat-hf', vocab_size=32000, model_max_length=1000000000000000019884624838656, is_fast=False, padding_side='right', truncation_side='right', special_tokens={'bos_token': '<s>', 'eos_token': '</s>', 'unk_token': '<unk>', 'pad_token': '▁***'}, clean_up_tokenization_spaces=False),  added_tokens_decoder={\n",
              "\t0: AddedToken(\"<unk>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
              "\t1: AddedToken(\"<s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
              "\t2: AddedToken(\"</s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
              "\t32000: AddedToken(\"<pad>\", rstrip=False, lstrip=False, single_word=False, normalized=True, special=False),\n",
              "}"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Util function for generate result\n",
        "def test_model(model, tokenizer, prompt, max_new_tokens=400, top_p=0.9, temperature=0.7):\n",
        "    \"\"\"\n",
        "    Generates text using a provided model and tokenizer.\n",
        "\n",
        "    Args:\n",
        "    - model: The language model to use for generation.\n",
        "    - tokenizer: The tokenizer associated with the model.\n",
        "    - prompt: The prompt to feed to the model.\n",
        "    - max_new_tokens: The maximum number of new tokens to generate. Default is 400.\n",
        "    - top_p: Nucleus sampling's cumulative probability cutoff. Default is 0.9.\n",
        "    - temperature: Controls randomness in generation. Lower values make text less random. Default is 0.7.\n",
        "\n",
        "    Returns:\n",
        "    A string containing the generated text.\n",
        "    \"\"\"\n",
        "\n",
        "    # Tokenize the prompt\n",
        "    tmp_eos = tokenizer.add_eos_token\n",
        "    tokenizer.add_eos_token = False\n",
        "    input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids.cuda()\n",
        "    tokenizer.add_eos_token = tmp_eos\n",
        "\n",
        "    # Generate the output\n",
        "    outputs = model.generate(\n",
        "        input_ids=input_ids,\n",
        "        max_new_tokens=max_new_tokens,\n",
        "        top_p=top_p,\n",
        "        do_sample=True,\n",
        "        temperature=temperature,\n",
        "        eos_token_id=tokenizer.eos_token_id\n",
        "    )\n",
        "\n",
        "    # Decode and clean up the output\n",
        "    generated_output = tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]\n",
        "    print(f\"Prompt:\\n{prompt}\\n\")\n",
        "    print(f\"Generated output:\\n{generated_output}\")\n",
        "    return"
      ],
      "metadata": {
        "id": "wgt-EMXDPhWz"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Test original model"
      ],
      "metadata": {
        "id": "vhjnblaUl11m"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = f\"\"\"[INST] <<SYS>>你是一个唐诗助手,帮助用户写一首对应要求的唐诗<</SYS>>\n",
        "作者:駱賓王\n",
        "标签:思念;七言律诗;秋天;咏物\n",
        "[/INST]\n",
        "\"\"\"\n",
        "\n",
        "test_model(model, tokenizer, prompt)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "06kw3iv9PGMz",
        "outputId": "c7ecfd05-34e0-4751-87ac-54e24bf21da9"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py:1518: UserWarning: You have modified the pretrained model configuration to control generation. This is a deprecated strategy to control generation and will be removed soon, in a future version. Please use and modify the model generation configuration (see https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prompt:\n",
            "[INST] <<SYS>>你是一个唐诗助手,帮助用户写一首对应要求的唐诗<</SYS>>\n",
            "作者:駱賓王\n",
            "标签:思念;七言律诗;秋天;咏物\n",
            "[/INST]\n",
            "\n",
            "\n",
            "Generated output:\n",
            "凛冬风雨晨，思念萧瑟潇湿。\n",
            "Autumn winds howl through the night,\n",
            "Memories of you, my heart's delight.\n",
            "\n",
            "In the morning dew, I wander alone,\n",
            "Missing you, my dear, like a forgotten tone.\n",
            "The trees stand tall, their leaves now gold,\n",
            "But without you, my heart's cold.\n",
            "\n",
            "Seven lines, a perfect rhyme,\n",
            "A poem of longing, a heart's crime.\n",
            "In autumn's hue, my thoughts are cast,\n",
            "Longing for you, my love, at last.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Initilize training parameters"
      ],
      "metadata": {
        "id": "oWAftXuJl5ib"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from peft import LoraConfig, PeftModel, PeftConfig, prepare_model_for_kbit_training, get_peft_model\n",
        "from transformers import TrainingArguments\n",
        "from trl import SFTTrainer\n",
        "\n",
        "peft_config = LoraConfig(\n",
        "    r=lora_r,\n",
        "    lora_alpha=lora_alpha,\n",
        "    lora_dropout=lora_dropout,\n",
        "    target_modules=lora_target_modules,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\",\n",
        ")\n",
        "\n",
        "args = TrainingArguments(\n",
        "    output_dir=output_dir,\n",
        "    num_train_epochs=num_train_epochs,\n",
        "    per_device_train_batch_size=per_device_train_batch_size,\n",
        "    gradient_accumulation_steps=gradient_accumulation_steps,\n",
        "    optim=optim,\n",
        "    logging_steps=logging_steps,\n",
        "    learning_rate=learning_rate,\n",
        "    fp16=fp16,\n",
        "    bf16=bf16,\n",
        "    max_grad_norm=max_grad_norm,\n",
        "    warmup_ratio=warmup_ratio,\n",
        "    weight_decay=weight_decay,\n",
        "    lr_scheduler_type=lr_scheduler_type,\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xV7233N6QabJ",
        "outputId": "4eee2811-b9f1-4f86-d4f6-7b5dca2b5bfa"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/trl/trainer/ppo_config.py:141: UserWarning: The `optimize_cuda_cache` arguement will be deprecated soon, please use `optimize_device_cache` instead.\n",
            "  warnings.warn(\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model = prepare_model_for_kbit_training(model)\n",
        "model = get_peft_model(model, peft_config)\n",
        "model.print_trainable_parameters()\n",
        "trainer = SFTTrainer(\n",
        "    model=model,\n",
        "    train_dataset=dataset,\n",
        "    peft_config=peft_config,\n",
        "    max_seq_length=max_seq_length,\n",
        "    tokenizer=tokenizer,\n",
        "    packing=packing,\n",
        "    dataset_text_field=\"text\",\n",
        "    args=args,\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 66,
          "referenced_widgets": [
            "a593780524b4420da3cc4e907bcc8b07",
            "ce9323cc706e4752bc9474c23daba43a",
            "17a2a78957654181a21d7bf86d6caaaf",
            "09299c4279be41ea932756b77380ca58",
            "c715443800cd4a129ad4741116363a2f",
            "af0147dc68a245d7ad01c088dc21ab7c",
            "265014e8078b4484ac590d5ea10c0804",
            "b97668c0895346f29d6f75a0e77175e9",
            "cbe5954851214e2a91ce943adf0c64d5",
            "ce3dfb1882e646d0ac247eda73cd2565",
            "2ec5b9ac76a54e83b2fb9a4ec319acd1"
          ]
        },
        "id": "OnQ29Q0mQumA",
        "outputId": "d85b7133-86f6-48f9-8769-93e8acb66a42"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "trainable params: 159,907,840 || all params: 6,898,323,456 || trainable%: 2.3180681656919973\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Map:   0%|          | 0/366 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "a593780524b4420da3cc4e907bcc8b07"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "trainer.train()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "Mmwhksu3Q2Vk",
        "outputId": "3fce2afb-95d9-42f2-8bab-0b3df013d665"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='366' max='366' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [366/366 28:25, Epoch 3/3]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Step</th>\n",
              "      <th>Training Loss</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>10</td>\n",
              "      <td>3.678500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>20</td>\n",
              "      <td>2.451900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>30</td>\n",
              "      <td>1.975600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>40</td>\n",
              "      <td>1.828100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>50</td>\n",
              "      <td>1.835700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>60</td>\n",
              "      <td>1.647800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>70</td>\n",
              "      <td>1.674100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>80</td>\n",
              "      <td>1.585300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>90</td>\n",
              "      <td>1.609800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>100</td>\n",
              "      <td>1.607500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>110</td>\n",
              "      <td>1.515700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>120</td>\n",
              "      <td>1.779800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>130</td>\n",
              "      <td>1.519700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>140</td>\n",
              "      <td>1.398500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>150</td>\n",
              "      <td>1.386600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>160</td>\n",
              "      <td>1.472500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>170</td>\n",
              "      <td>1.214000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>180</td>\n",
              "      <td>1.281200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>190</td>\n",
              "      <td>1.229200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>200</td>\n",
              "      <td>1.363700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>210</td>\n",
              "      <td>1.355900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>220</td>\n",
              "      <td>1.427400</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>230</td>\n",
              "      <td>1.246300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>240</td>\n",
              "      <td>1.252500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>250</td>\n",
              "      <td>1.116200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>260</td>\n",
              "      <td>1.104100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>270</td>\n",
              "      <td>1.226600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>280</td>\n",
              "      <td>1.025700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>290</td>\n",
              "      <td>1.078400</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>300</td>\n",
              "      <td>1.247000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>310</td>\n",
              "      <td>1.087000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>320</td>\n",
              "      <td>0.991000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>330</td>\n",
              "      <td>1.024000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>340</td>\n",
              "      <td>1.040000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>350</td>\n",
              "      <td>0.903500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>360</td>\n",
              "      <td>0.902500</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TrainOutput(global_step=366, training_loss=1.4362254742064762, metrics={'train_runtime': 1709.92, 'train_samples_per_second': 0.642, 'train_steps_per_second': 0.214, 'total_flos': 1.0621756074811392e+16, 'train_loss': 1.4362254742064762, 'epoch': 3.0})"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# save model\n",
        "trainer.save_model(save_dir)"
      ],
      "metadata": {
        "id": "kyL4xLgWRA2I"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Empty VRAM\n",
        "del model\n",
        "del trainer\n",
        "import gc\n",
        "gc.collect()\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "sEopR7VxRgV-"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# load saved model\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.bfloat16\n",
        ")\n",
        "\n",
        "config = PeftConfig.from_pretrained(save_dir)\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    config.base_model_name_or_path,\n",
        "    return_dict=True,\n",
        "    quantization_config=bnb_config,\n",
        "    device_map=\"auto\",\n",
        "    trust_remote_code=True\n",
        ")\n",
        "model = PeftModel.from_pretrained(model, save_dir)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 138,
          "referenced_widgets": [
            "d62f31bc09054d72846e601dcdb167f2",
            "78a776d88c944c558aa8654e08ccae39",
            "eafe0b3a7b324bdbafc6d28da64885dc",
            "7b6253d809ea4b3588800f87b2ddd641",
            "9998068daf81422aae5b0e2e3223b11a",
            "4551d7c16dd540eca3d8fe819e14b221",
            "62da5e0db3dd47c786d56c7d7828f38b",
            "ec8913ad1253489bafcfdba52fbcb528",
            "a2b8a05d722c40058cc7f6868d9f086b",
            "7f14dd60d1aa45fabecf5d98168f7be4",
            "4fc60b58df5648cba3f9a8fc865bfeb8"
          ]
        },
        "id": "69N1fSJaRfCR",
        "outputId": "9e25643d-7783-4736-a992-08006a11a786"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "d62f31bc09054d72846e601dcdb167f2"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/transformers/generation/configuration_utils.py:389: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.9` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`. This was detected when initializing the generation config instance, which means the corresponding file may hold incorrect parameterization and should be fixed.\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.10/dist-packages/transformers/generation/configuration_utils.py:394: UserWarning: `do_sample` is set to `False`. However, `top_p` is set to `0.6` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_p`. This was detected when initializing the generation config instance, which means the corresponding file may hold incorrect parameterization and should be fixed.\n",
            "  warnings.warn(\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Check finetune result"
      ],
      "metadata": {
        "id": "q1jw7Q_kmD2J"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = f\"\"\"<s>[INST] <<SYS>>你是一个唐诗助手,帮助用户写一首对应要求的唐诗<</SYS>>\n",
        "作者:李白\n",
        "标签:乐府;赞美;近代曲辞\n",
        "[/INST]\"\"\"\n",
        "test_model(model, tokenizer, prompt)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "41Nb3TdAat1f",
        "outputId": "f7836816-ea1a-4487-b189-0c962843e641"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prompt:\n",
            "<s>[INST] <<SYS>>你是一个唐诗助手,帮助用户写一首对应要求的唐诗<</SYS>>\n",
            "作者:李白\n",
            "标签:乐府;赞美;近代曲辞\n",
            "[/INST]\n",
            "\n",
            "Generated output:\n",
            "辭 清平調 三 \n",
            "樂聖樂天下，妾心中欲妝。\n",
            "誰調彈劒輕，一把琴聽行。\n",
            "鳴琴吞聲賦，感懷滋潤勞。\n",
            "欲妝妾心怯，欲慰心欲殘。\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = f\"\"\"[INST] <<SYS>>你是一个唐诗助手,帮助用户写一首对应要求的唐诗<</SYS>>\n",
        "作者:李商隱\n",
        "标签:黄河;咏物;抒情;鼓吹曲辞;乐府;咏物诗\n",
        "[/INST]\n",
        "\"\"\"\n",
        "test_model(model, tokenizer, prompt)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1z1f9m7oazvr",
        "outputId": "dbe29e76-de94-466f-9d5c-511397fe0e03"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prompt:\n",
            "[INST] <<SYS>>你是一个唐诗助手,帮助用户写一首对应要求的唐诗<</SYS>>\n",
            "作者:李商隱\n",
            "标签:黄河;咏物;抒情;鼓吹曲辞;乐府;咏物诗\n",
            "[/INST]\n",
            "\n",
            "\n",
            "Generated output:\n",
            "鼓吹曲辭 黄河 \n",
            "黃河出山東，兩岸胡時雨。\n",
            "輕舟橫漠漠，閑坐覽夕陽。\n",
            "兩岸廢舊村，一片荒涼草。\n",
            "黃河又何處，千里寒雲滿。\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = f\"\"\"<s>[INST] <<SYS>>你是一个唐诗助手,帮助用户写一首对应要求的唐诗<</SYS>>\n",
        "作者:杜甫\n",
        "标签:乐府;赞美;近代曲辞\n",
        "[/INST]\n",
        "\"\"\"\n",
        "test_model(model, tokenizer, prompt)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Tux45SJ8az2U",
        "outputId": "83508567-f4b5-4631-e4f5-b8734ac972ab"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prompt:\n",
            "<s>[INST] <<SYS>>你是一个唐诗助手,帮助用户写一首对应要求的唐诗<</SYS>>\n",
            "作者:杜甫\n",
            "标签:乐府;赞美;近代曲辞\n",
            "[/INST]\n",
            "\n",
            "\n",
            "Generated output:\n",
            "辭 清平調 一 \n",
            "清平調，欲往舟卒，終於晚泊。\n",
            "獨立鳳雛，萬里長風。\n",
            "昔日青春，不可復來。\n",
            "潮自復興，海燕自遁散。\n"
          ]
        }
      ]
    }
  ]
}